import numpy as np  # linear algebra
import pandas as pd  # data processing, CSV file I/O (e.g. pd.read_csv)
import matplotlib.pyplot as plt

from sklearn.model_selection import train_test_split
from sklearn import preprocessing
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

from datasets import MyDataset

from network import NerualNet_yieldRate



def test(dataPath, modelSavePath):
    print('\ntraining: ', dataPath)

    dataset = pd.read_csv(dataPath)
    # X = dataset.iloc[:, :-1]
    # y = dataset.iloc[:, dataset.shape[1] - 1:dataset.shape[1]]
    X = dataset.iloc[:, :-1]
    y = dataset['L']

    # normalize
    # X_scaled = preprocessing.StandardScaler().fit_transform(X)
    # X_scaled = preprocessing.MinMaxScaler().fit_transform(X)
    # X_scaled = preprocessing.MaxAbsScaler().fit_transform(X)
    X_scaled = preprocessing.Normalizer().fit_transform(X)

    # X_scaled = None
    # if dataPath=='../../data/ingotRate.csv':
    #     X_scaled = preprocessing.MaxAbsScaler().fit_transform(X)
    # elif dataPath=='../../data/yieldRate.csv':
    #     X_scaled = preprocessing.StandardScaler().fit_transform(X)

    # y_scaled = preprocessing.Normalizer().fit_transform(y)

    # Splitting the dataset into the Training set and Test set
    X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.2, random_state=0)

    # print('\n')
    # print('\ntrain image shape : ', X_train.shape)
    # print('train label shape : ', y_train.shape)
    # print('test image shape  : ', X_test.shape)
    # print('valid label image : ', y_test.shape)

    test_dataset = MyDataset(X=X_test, y=y_test)

    test_loader = DataLoader(dataset=test_dataset, batch_size=128, shuffle=False)
    #
    # model = NerualNet_yieldRate()
    # print('\nmodel: ')
    # print(model)
    #
    # optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    # loss_fn = nn.MSELoss()
    #
    # test_losses = {'MSE': [], 'MAE': [], 'R2': []}
    # with torch.no_grad():
    #     for i, (data, label) in enumerate(test_loader):
    #         data = data.float()
    #         label = label.float()
    #
    #         output = model(data)
    #         loss = loss_fn(output, label)
    #
    #         MAE = mean_absolute_error(output.detach().numpy(), label.detach().numpy())
    #         R2_Score = r2_score(output.detach().numpy(), label.detach().numpy())
    #         test_losses['MSE'].append(loss.item())
    #         test_losses['MAE'].append(MAE)
    #         test_losses['R2'].append(R2_Score)

    model = torch.load(modelSavePath)
    # output = model(X_test)
    # print(output)


    loss_fn = nn.MSELoss()

    model.eval()
    test_preds = torch.LongTensor()

    test_losses = {'MSE': [], 'MAE': [], 'R2': []}
    y_pred=None
    for i, (data, label) in enumerate(test_loader):
        data=data.float()
        outputs = model(data)
        # print(outputs)
        y_pred=outputs.detach().numpy()

        loss = loss_fn(outputs, label)

        MAE = mean_absolute_error(outputs.detach().numpy(), label.detach().numpy())
        R2_Score = r2_score(outputs.detach().numpy(), label.detach().numpy())
        test_losses['MSE'].append(loss.item())
        test_losses['MAE'].append(MAE)
        test_losses['R2'].append(R2_Score)


    plt.plot(y_test.values, label='y_test')
    plt.plot(y_pred, label='y_pred')
    plt.title('ours_LFDS', fontsize='large', fontweight = 'bold')
    plt.legend()
    plt.show()

    print('\ntest: ')
    print("MSE: ", test_losses['MSE'][0])
    print("MAE: ", test_losses['MAE'][0])
    print("R2_Score: ", test_losses['R2'][0])



if __name__ == '__main__':
    test('../../data/ingotRate.csv', './ingotRatePrediction.pkl')